"""
期货分析智能体模块 - 提供期货分析功能
"""

from typing import Dict, TypedDict, List, Union, Any, Callable
from langgraph.graph import StateGraph, END
import json
import datetime
import logging
import pandas as pd

from FDAS.agents.base import BaseAgent, BaseAgentState
from FDAS.data import get_asset
from FDAS.data.assets.base import ASSET_TYPE_FUTURE
from FDAS.agents.workflows.prompts import (
    get_query_parse_prompt,
    get_specific_analysis_prompt,
    get_synthesis_prompt
)

# 自定义JSON编码器，处理日期类型
class DateTimeEncoder(json.JSONEncoder):
    def default(self, obj):
        if isinstance(obj, (datetime.date, datetime.datetime)):
            return obj.isoformat()
        return super(DateTimeEncoder, self).default(obj)

# 定义期货智能体状态
class FutureAgentState(BaseAgentState):
    symbol: str = ""
    analysis_type: str = ""
    future_data: Dict = {}
    # 各子分析结果
    basic_info_analysis: str = ""
    price_analysis: str = ""
    technical_analysis: str = ""

class FutureAnalysisAgent(BaseAgent):
    """期货分析智能体"""
    
    def _build_workflow(self) -> StateGraph:
        """构建智能体工作流"""
        workflow = StateGraph(FutureAgentState)
        
        # 添加节点
        workflow.add_node("parse_query", self._parse_query)
        workflow.add_node("retrieve_data", self._retrieve_data)
        
        # 添加各个分析节点
        workflow.add_node("analyze_basic_info", self._analyze_basic_info)
        workflow.add_node("analyze_price", self._analyze_price)
        workflow.add_node("analyze_technical", self._analyze_technical)
        
        # 路由节点
        workflow.add_node("route_next", self._route_to_next_analysis)
        
        # 综合分析节点
        workflow.add_node("synthesize_analysis", self._synthesize_analysis)
        
        # 设置节点连接关系 - 基本工作流
        workflow.set_entry_point("parse_query")
        workflow.add_edge("parse_query", "retrieve_data")
        
        # 从数据获取到各个分析节点的连接
        workflow.add_edge("retrieve_data", "analyze_basic_info")
        
        # 从各个分析节点到路由节点的连接
        workflow.add_edge("analyze_basic_info", "route_next")
        workflow.add_edge("analyze_price", "route_next")
        workflow.add_edge("analyze_technical", "route_next")
        
        # 使用条件边缘路由到下一个节点
        workflow.add_conditional_edges(
            "route_next",
            self._decide_next_node,
            {
                "analyze_price": "analyze_price",
                "analyze_technical": "analyze_technical",
                "synthesize_analysis": "synthesize_analysis"
            }
        )
        
        # 从综合分析到结束的连接
        workflow.add_edge("synthesize_analysis", END)
        
        return workflow
    
    def _decide_next_node(self, state: FutureAgentState) -> str:
        """确定下一个要执行的节点（为条件边缘决策函数）"""
        # 检查数据和分析类型
        future_data = state.get('future_data', {})
        analysis_type = state.get('analysis_type', '')
        
        # 价格分析
        if analysis_type in ['价格分析', '综合分析'] and 'history_data' in future_data and 'price_analysis' not in state:
            return "analyze_price"
        
        # 技术分析
        if analysis_type in ['技术分析', '综合分析'] and 'history_data' in future_data and 'technical_analysis' not in state:
            return "analyze_technical"
        
        # 所有必要分析都已完成，进行综合分析
        return "synthesize_analysis"
    
    def _route_to_analysis_nodes(self, state: FutureAgentState) -> str:
        """根据分析类型和数据可用性路由到适当的分析节点"""
        # 首先始终进行基本信息分析
        return "analyze_basic_info"
    
    def _route_to_next_analysis(self, state: FutureAgentState) -> Dict:
        """检查并决定下一个分析节点 - 此函数仅用于边缘连接，不决定实际路由"""
        # 此函数作为节点函数，必须返回字典以更新状态
        # 在使用条件边缘时，节点的输出不决定路由，而是由条件边缘的决策函数决定
        return state
    
    def _parse_query(self, state: FutureAgentState) -> Dict:
        """解析用户查询并提取期货代码和分析类型"""
        try:
            prompt = get_query_parse_prompt("future")
            response = self.llm.invoke(prompt.format(user_query=state['user_query']))
            
            # 解析JSON
            result = self._clean_json_response(response.content)
            
            # 数据校验
            result['symbol'] = str(result.get('symbol', 'CL')).strip()
            analysis_type = str(result.get('analysis_type', '技术分析')).lower()
            result['analysis_type'] = {
                'price': '价格分析',
                'technical': '技术分析',
                'comprehensive': '综合分析',
                '价格': '价格分析',
                '技术': '技术分析',
                '综合': '综合分析'
            }.get(analysis_type, '技术分析')
            
            return {
                "symbol": result['symbol'],
                "analysis_type": result['analysis_type'],
                "parsed_query": json.dumps(result, ensure_ascii=False, cls=DateTimeEncoder),
                "user_query": state['user_query']
            }
                
        except Exception as e:
            self.logger.error(f"查询解析失败: {str(e)}")
            # 默认返回
            return {
                "symbol": "CL",
                "analysis_type": "技术分析",
                "parsed_query": f"解析失败：{str(e)}",
                "user_query": state['user_query']
            }
    
    def _retrieve_data(self, state: FutureAgentState) -> Dict:
        """获取相关数据"""
        try:
            symbol = state['symbol']
            analysis_type = state['analysis_type']
            
            # 获取期货对象
            future = get_asset(symbol, ASSET_TYPE_FUTURE)
            
            # 基础数据 - 所有分析类型都需要
            basic_info = future.basic_info
            
            # 根据分析类型获取不同数据
            if analysis_type == '价格分析':
                data = {
                    "basic_info": basic_info,
                    "history_data": future.get_history(limit=60).to_dict('records')
                }
            elif analysis_type == '技术分析':
                data = {
                    "basic_info": basic_info,
                    "history_data": future.get_history(limit=120).to_dict('records')
                }
            elif analysis_type == '综合分析':
                data = {
                    "basic_info": basic_info,
                    "history_data": future.get_history(limit=120).to_dict('records'),
                    "open_interest": future.get_open_interest().to_dict('records')
                }
            else:
                data = {"error": f"暂不支持的分析类型: {analysis_type}"}
            
            return {
                **state,  # 保留现有状态
                "future_data": data
            }
                
        except Exception as e:
            self.logger.error(f"数据获取失败: {str(e)}")
            return {
                **state,
                "future_data": {"error": f"数据获取失败: {str(e)}"}
            }
    
    def _analyze_basic_info(self, state: FutureAgentState) -> Dict:
        """分析基本信息"""
        try:
            # 获取相关数据
            future_data = state.get('future_data', {})
            basic_info = future_data.get('basic_info', {})
            
            # 检查数据
            if not basic_info:
                return {**state, "basic_info_analysis": "未获取到基本信息数据"}
            
            # 构建提示
            prompt = get_specific_analysis_prompt('basic_info', 'future').format(
                basic_info=json.dumps(basic_info, ensure_ascii=False, cls=DateTimeEncoder)
            )
            
            # 调用大模型
            response = self.llm.invoke(prompt)
            
            # 返回结果
            return {**state, "basic_info_analysis": response.content}
                
        except Exception as e:
            self.logger.error(f"基本信息分析失败: {str(e)}")
            return {**state, "basic_info_analysis": f"分析失败: {str(e)}"}
    
    def _analyze_price(self, state: FutureAgentState) -> Dict:
        """分析价格数据"""
        try:
            # 获取相关数据
            future_data = state.get('future_data', {})
            basic_info = future_data.get('basic_info', {})
            history_data = future_data.get('history_data', [])
            
            # 检查数据
            if not history_data:
                return {**state, "price_analysis": "未获取到历史价格数据"}
            
            # 限制数据量 - 只取前30条
            limited_history = history_data[:30] if len(history_data) > 30 else history_data
            
            # 构建提示
            prompt = get_specific_analysis_prompt('price', 'future').format(
                basic_info=json.dumps(basic_info, ensure_ascii=False, cls=DateTimeEncoder),
                history_data=json.dumps(limited_history, ensure_ascii=False, cls=DateTimeEncoder)
            )
            
            # 调用大模型
            response = self.llm.invoke(prompt)
            
            # 返回结果
            return {**state, "price_analysis": response.content}
                
        except Exception as e:
            self.logger.error(f"价格分析失败: {str(e)}")
            return {**state, "price_analysis": f"分析失败: {str(e)}"}
    
    def _analyze_technical(self, state: FutureAgentState) -> Dict:
        """分析技术指标"""
        try:
            # 获取相关数据
            future_data = state.get('future_data', {})
            basic_info = future_data.get('basic_info', {})
            history_data = future_data.get('history_data', [])
            
            # 检查数据
            if not history_data:
                return {**state, "technical_analysis": "未获取到历史价格数据，无法计算技术指标"}
            
            # 限制数据量 - 只取前30条用于计算技术指标
            limited_history = history_data[:30] if len(history_data) > 30 else history_data
            
            # 构建提示
            prompt = get_specific_analysis_prompt('technical', 'future').format(
                basic_info=json.dumps(basic_info, ensure_ascii=False, cls=DateTimeEncoder),
                history_data=json.dumps(limited_history, ensure_ascii=False, cls=DateTimeEncoder)
            )
            
            # 调用大模型
            response = self.llm.invoke(prompt)
            
            # 返回结果
            return {**state, "technical_analysis": response.content}
                
        except Exception as e:
            self.logger.error(f"技术指标分析失败: {str(e)}")
            return {**state, "technical_analysis": f"分析失败: {str(e)}"}
    
    def _synthesize_analysis(self, state: FutureAgentState) -> Dict:
        """综合各个分析，生成最终报告"""
        try:
            # 收集所有分析总结
            summaries = {}
            for key in [
                "basic_info_analysis", "price_analysis", "technical_analysis"
            ]:
                if key in state and state.get(key):
                    summaries[key] = state.get(key)
            
            # 构建提示
            prompt = get_synthesis_prompt('future').format(
                user_query=state.get('user_query', ''),
                symbol=state.get('symbol', ''),
                summaries=json.dumps(summaries, ensure_ascii=False)
            )
            
            # 调用大模型
            response = self.llm.invoke(prompt)
            
            # 返回结果
            return {**state, "analysis_result": response.content}
                
        except Exception as e:
            self.logger.error(f"综合分析失败: {str(e)}")
            return {**state, "analysis_result": self._format_error_response(str(e))}
    
    def analyze(self, query: str) -> str:
        """执行完整分析流程"""
        try:
            result = self.agent.invoke({"user_query": query})
            return result['analysis_result']
        except Exception as e:
            self.logger.error(f"分析执行失败: {str(e)}")
            return self._format_error_response(str(e))